package demo.datastram;

import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.dag.Transformation;
import org.apache.flink.api.java.Utils;
import org.apache.flink.api.java.typeutils.TypeExtractor;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
import org.apache.flink.streaming.api.operators.SimpleOperatorFactory;
import org.apache.flink.streaming.api.operators.StreamMap;
import org.apache.flink.streaming.api.operators.StreamOperatorFactory;
import org.apache.flink.streaming.api.transformations.OneInputTransformation;

public class DataStreamDemo<T> {

    protected final StreamExecutionEnvironment environment;
    protected final Transformation transformation;

    public DataStreamDemo(StreamExecutionEnvironment environment, Transformation transformation) {
        this.environment = environment;
        this.transformation = transformation;
    }

    TypeInformation<T> getType() {
        return  transformation.getOutputType();
    }

    public <R> SingleOutputStreamOperatorDemo<R> map(MapFunction<T, R> mapper) {
        TypeInformation<R> outType = TypeExtractor.getMapReturnTypes(mapper, getType(), Utils.getCallLocationName(), true);
        return map(mapper, outType);
    }

    public <R> SingleOutputStreamOperatorDemo<R> map(MapFunction<T, R> mapper, TypeInformation<R> outputType) {
        return transform("Map", outputType, new StreamMap<>(mapper));
    }

    public <R> SingleOutputStreamOperatorDemo<R> transform (
            String operatorName,
            TypeInformation<R> outTypeInfo,
            OneInputStreamOperator<T, R> operator
    ) {
        return doTransform(operatorName, outTypeInfo, SimpleOperatorFactory.of(operator));
    }

    public <R> SingleOutputStreamOperatorDemo<R> doTransform (
            String operatorName,
            TypeInformation<R> outTypeInfo,
            StreamOperatorFactory<R> operatorFactor
    ) {
        transformation.getOutputType();
        OneInputTransformation resultTransform = new OneInputTransformation<>(
                this.transformation,
                operatorName,
                operatorFactor,
                outTypeInfo,
                environment.getParallelism()
        );
        SingleOutputStreamOperatorDemo<R> returnStream = new SingleOutputStreamOperatorDemo(environment, resultTransform);

        getExecutionEnvironment().addOperator(resultTransform);
        return returnStream;
    }

    public StreamExecutionEnvironment getExecutionEnvironment() {
        return environment;
    }

}
